#include <stdio.h>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv/cv.h>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/ml/ml.hpp>

/**
 * @brief The method that helps user to interact with programme.
 * @param input
 * @return The value of key input.
 * @author 1160300719 殷浩然
 */
int showInformation(std::string input){
    cv::Mat image = cv::Mat::zeros(cv::Size(640, 240), CV_8UC3);
    image.setTo(cv::Scalar(205, 250, 255));
    std::string text = input;
    int font_face = cv::FONT_HERSHEY_COMPLEX;
    double font_scale = 0.9;
    int thickness = 1;
    int baseline;
    cv::Size text_size = cv::getTextSize(text, font_face, font_scale, thickness, &baseline);
    cv::Point origin;
    origin.x = image.cols / 2 - text_size.width / 2;
    origin.y = image.rows / 2 + text_size.height / 2;
    putText(image, text, origin, font_face, font_scale, cv::Scalar(0, 0, 0), thickness, 8, 0);
    cv::imshow("image", image);
    return cv::waitKey();
}

/**
 * @brief The method that extract the information of the number image.
 * @param src_img
 * @author 1160300719 殷浩然
 */
void getNumber(cv::Mat src_img){
    cv::Mat gray_img, blur_img, dst_img;
    cv::GaussianBlur(src_img, blur_img, cv::Size(5, 5), 0);
    cv::Canny(blur_img, gray_img, 1, 500);
    cv::Mat ROI_img = gray_img(cv::Rect(190, 75, 100, 100));
    cv::imwrite("../picture/gray_img.jpg", gray_img);
    cv::imwrite("../picture/ROI.jpg", ROI_img);
    std::vector<std::vector<cv::Point> > contours;
    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(ROI_img, contours, hierarchy, CV_RETR_CCOMP, CV_CHAIN_APPROX_SIMPLE);
    for ( size_t i = 0; i < contours.size(); i++ ){
        cv::drawContours( ROI_img, contours, i, cv::Scalar(255), CV_FILLED, 8, std::vector<cv::Vec4i>(), 0, cv::Point() );
    }
    cv::resize(ROI_img, ROI_img, cv::Size(), 0.2, 0.2);
    cv::imwrite("../picture/result.jpg", ROI_img);
}

/**
 * @brief The method that read the data(/data) to intialize the training of SVM.
 * @param trainingImages
 * @param trainingLabels
 * @author 1160300719 殷浩然
 */
void getTrainData(cv::Mat& trainingImages, std::vector<int>& trainingLabels){
    for(unsigned int k=0; k<10; k++){
        std::vector<std::string> files;
        for(unsigned int i=0; i<500; i++){
            char ch[20];
            sprintf(ch, "../data/%d/%d.jpg", k, i);
            files.push_back(ch);
        }
        int number = (int)files.size();
        for (unsigned int i=0; i<number; i++)
        {
            cv::Mat SrcImage = cv::imread(files[i].c_str());
            SrcImage = SrcImage.reshape(1, 1);
            trainingImages.push_back(SrcImage);
            trainingLabels.push_back(k);
        }
    }
}

/**
 * @brief The method that save the result that gets from SVM.
 * @author 1160300719 殷浩然
 */
void train(){
    cv::Mat classes;
    cv::Mat trainingData;
    cv::Mat trainingImages;
    std::vector<int> trainingLabels;
    getTrainData(trainingImages, trainingLabels);
    cv::Mat(trainingImages).copyTo(trainingData);
    trainingData.convertTo(trainingData, CV_32FC1);
    cv::Mat(trainingLabels).copyTo(classes);
    cv::Ptr<cv::ml::SVM> svm = cv::ml::SVM::create();
    svm->setType(100);
    svm->setKernel(0);
    svm->setTermCriteria(cv::TermCriteria(cv::TermCriteria::MAX_ITER, 1000, 0.01));
    svm->setDegree(0);
    svm->setGamma(1);
    svm->setCoef0(0);
    svm->setC(1);
    svm->setNu(0);
    svm->setP(0);
    svm->train(trainingData, 0, classes);
    svm->save("../xml/svm.xml");
}

/**
 * @brief Follow the logic and combine the method above.
 * @return The running state of programme.
 * @author 1160300719 殷浩然
 */
int main()
{
    int first_choice = showInformation("Do you want to train SVM first?(y/n)");
    if(first_choice == 121){
        train();
        showInformation("Success! Press any key to continue.");
    }
    cv::Mat src_img = cv::imread("../picture/1.png");
    getNumber(src_img);
    cv::Mat dst_img = cv::imread("../picture/result.jpg");
    dst_img = dst_img.reshape(1, 1);
    dst_img.convertTo(dst_img, CV_32FC1);
    cv::Ptr<cv::ml::SVM> svm = cv::Algorithm::load<cv::ml::SVM>("../xml/svm.xml");
    float response = svm->predict(dst_img);
    char ch[40];
    sprintf(ch, "Predict result: %d", (int)response);
    showInformation(ch);
    return  0;
}